package com.ruoyi.stock2.fasttext;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.StringReader;
import java.io.UnsupportedEncodingException;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.th.ThaiAnalyzer;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.wltea.analyzer.core.IKSegmenter;
import org.wltea.analyzer.core.Lexeme;

public class TrainModel {

	private static String tab = " ";

	public static void main(String[] args) {
		args = new String[] { "D:/ChnSentiCorp情感分析酒店评论", "D:/model", "subject", "cn" };

//		args = new String[] { "D:/thai-data", "D:/model", "subject", "th" };
		try {
			// 语料路径parameter1
			// 输出训练模型parameter2
			// 模型主题parameter3
			// 英文或泰文parameter4 en英文 th泰文 cn中文
//			String[] typeArray = new String[] { "car", "game", "sports" };
			String[] typeArray = new String[] { "negative", "positive" };
			String tmpCorpus = args[1] + File.separator + args[2] + ".txt";
			// 语料文件合并打标记
			TrainModel.processCorpus(typeArray, args[0], tmpCorpus, args[3]);
			// 训练模型
			TrainModel.trainModel(args[1], args[2], tmpCorpus);

		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public static String readToString(String fileName) {
		String encoding = "utf-8";
		File file = new File(fileName);
		Long filelength = file.length();
		byte[] filecontent = new byte[filelength.intValue()];
		try {
			FileInputStream in = new FileInputStream(file);
			in.read(filecontent);
			in.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		try {
			return new String(filecontent, encoding);
		} catch (UnsupportedEncodingException e) {
			System.err.println("The OS does not support " + encoding);
			e.printStackTrace();
			return null;
		}
	}

	public static void writeTxtFile(String newStr, String path) throws IOException {
		byte[] sourceByte = newStr.getBytes("utf-8");
		if (null != sourceByte) {
			try {
				File file = new File(path); // 文件路径（路径+文件名）
				if (!file.exists()) { // 文件不存在则创建文件，先创建目录
					File dir = new File(file.getParent());
					dir.mkdirs();
					file.createNewFile();
				}
				FileOutputStream outStream = new FileOutputStream(file); // 文件输出流用于将数据写入文件
				outStream.write(sourceByte);
				outStream.close(); // 关闭文件输出流
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
	}

	public static void processCorpus(String[] typeArray, String args0, String tmpCorpus, String language)
			throws IOException {
		StringBuffer buffer = new StringBuffer("");
		Analyzer analyzer = new ThaiAnalyzer();
		for (String typeStr : typeArray) {
			File file = new File(args0 + File.separator + typeStr);
			if (file.exists()) {
				File[] files = file.listFiles();
				String content = "";
				for (File file2 : files) {
					content = TrainModel.readToString(file2.getAbsolutePath());
					if ("en".equals(language)) {
						buffer.append(content.replaceAll(System.getProperty("line.separator"), tab));
					}
					if ("th".equals(language)) {
						String thaiPro = TrainModel.thaiAnalzyer(analyzer, content);
						buffer.append(thaiPro.replaceAll(System.getProperty("line.separator"), tab));
					}
					if ("cn".equals(language)) {
						String thaiPro = TrainModel.cnAnalzyer(content);
						buffer.append(thaiPro.replaceAll(System.getProperty("line.separator"), tab));
					}
					buffer.append(tab + "__label__" + typeStr);
					buffer.append(System.getProperty("line.separator"));

				}
			}
		}
		TrainModel.writeTxtFile(buffer.toString(), tmpCorpus);
	}

	public static void trainModel(String args1, String args2, String tmpCorpus) throws IOException, Exception {
		String output = args1 + File.separator + args2 + ".model";
		String[] text = { "supervised", "-input", tmpCorpus, "-output", output, "-loss", "softmax", "-dim", "10", "-lr",
				"0.1", "-wordNgrams", "5", "-minCount", "1", "-bucket", "85000", "-epoch", "5", "-thread", "10" };
		Main op = new Main();
		op.train(text);
	}

	public static String thaiAnalzyer(Analyzer analyzer, String text) throws IOException {
		StringBuffer thaiSplit = new StringBuffer("");
		TokenStream tokenStream = analyzer.tokenStream("", text);
		tokenStream.addAttribute(CharTermAttribute.class);
		tokenStream.reset();
		while (tokenStream.incrementToken()) {
			CharTermAttribute charTermAttribute = tokenStream.getAttribute(CharTermAttribute.class);
			thaiSplit.append(new String(charTermAttribute.toString()) + tab);
		}
		tokenStream.close();
		return thaiSplit.toString();
	}

	public static String cnAnalzyer(String text) {
		StringBuffer cnSplit = new StringBuffer("");
		StringReader re = new StringReader(text.trim());
		IKSegmenter ik = new IKSegmenter(re, true);
		Lexeme lex = null;
		try {
			while ((lex = ik.next()) != null) {
				cnSplit.append(lex.getLexemeText() + tab);
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
		return cnSplit.toString();
	}
}
